package com.peng.rpc.loadBalance;

import com.alibaba.nacos.api.naming.pojo.Instance;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.time.Instant;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.function.Consumer;

/**
 * @author Wonderson
 * @email wpp.ds@foxmail.com
 * @date 2022/3/8 22:40
 */
public class ConsistentHashLoadBalance implements LoadBalance{

    private final static int REPLICA_NUMBER = 8;

    @Override
    public Instance select(List<Instance> instances) {
        ConsistentHashSelector selector = new ConsistentHashSelector(instances, REPLICA_NUMBER);
        return selector.select(new Date().toString());
    }

    static class ConsistentHashSelector {
        private final TreeMap<Long, Instance> virtualinstances;

        ConsistentHashSelector(List<Instance> instances, int replicaNumber) {
            this.virtualinstances = new TreeMap<>();
            for (Instance instance : instances) {
                for (int i = 0; i < replicaNumber / 4; i++) {
                    byte[] digest = md5(instance.getInstanceId() + i);
                    for (int h = 0; h < 4; h++) {
                        long m = hash(digest, h);
                        virtualinstances.put(m, instance);
                    }
                }
            }
        }

        static byte[] md5(String key) {
            MessageDigest md;
            try {
                md = MessageDigest.getInstance("MD5");
                byte[] bytes = key.getBytes(StandardCharsets.UTF_8);
                md.update(bytes);
            } catch (NoSuchAlgorithmException e) {
                throw new IllegalStateException(e.getMessage(), e);
            }

            return md.digest();
        }

        static long hash(byte[] digest, int idx) {
            return ((long) (digest[3 + idx * 4] & 255) << 24 | (long) (digest[2 + idx * 4] & 255) << 16 | (long) (digest[1 + idx * 4] & 255) << 8 | (long) (digest[idx * 4] & 255)) & 4294967295L;
        }

        public Instance select(String rpcServiceKey) {
            byte[] digest = md5(rpcServiceKey);
            return selectForKey(hash(digest, 0));
        }

        public Instance selectForKey(long hashCode) {
            Map.Entry<Long, Instance> entry = virtualinstances.tailMap(hashCode, true).firstEntry();

            if (entry == null) {
                //环形结构，若后面没有，就从环头部取得。
                entry = virtualinstances.firstEntry();
            }

            return entry.getValue();
        }
    }
}
